import numpy as np
import torch
import skimage as ski
from skimage import filters
from skimage import color
from skimage.util import invert
from skimage.exposure import match_histograms,equalize_hist,equalize_adapthist,is_low_contrast,rescale_intensity
from skimage.restoration import denoise_wavelet,denoise_tv_chambolle,denoise_bilateral
from einops import rearrange
def sk_sobel_filter(tensor):
    # 将 PyTorch tensor 转换为 NumPy ndarray
    img_np = tensor.numpy().astype(np.float64)
    
    # 应用 Sobel 滤波器
    edges = filters.sobel(img_np)
    
    # 将结果转换回 PyTorch tensor
    edges_tensor = torch.from_numpy(edges).float()
    
    return edges_tensor

def sk_laplace_filter(tensor):
    # 将 PyTorch tensor 转换为 NumPy ndarray
    img_np = tensor.numpy().astype(np.float64)
    
    # 应用 Laplacian 滤波器
    edges = filters.laplace(img_np)
    
    # 将结果转换回 PyTorch tensor
    edges_tensor = torch.from_numpy(edges).float()
    
    return edges_tensor

def new_int(JOIN,VIS,IR):
    # 将 PyTorch tensor 转换为 NumPy ndarray
    join = JOIN.numpy().astype(np.float64)
    vis = VIS.numpy().astype(np.float64)
    ir = IR.numpy().astype(np.float64)
    
    '''temp=invert(rescale_intensity(ir-vis))
    out=invert(temp-join)'''
    out=vis*ir*join
    # 应用 Laplacian 滤波器
    #edges = filters.laplace(out)
    
    # 将结果转换回 PyTorch tensor
    edges_tensor = torch.from_numpy(out).float()
    
    return edges_tensor
def new_int1(JOIN,VIS,IR):
    # 将 PyTorch tensor 转换为 NumPy ndarray
    join = JOIN.numpy().astype(np.float64)
    vis = VIS.numpy().astype(np.float64)
    ir = IR.numpy().astype(np.float64)
    
    '''temp=invert(rescale_intensity(ir-vis))
    out=invert(temp-join)'''
    out=vis*ir
    # 应用 Laplacian 滤波器
    #edges = filters.laplace(out)
    
    # 将结果转换回 PyTorch tensor
    edges_tensor = torch.from_numpy(out).float()
    
    return edges_tensor

def percentile_whitebalance(image, percentile_value):
    # 图片白平衡处理
    whitebalanced = rescale_intensity(image / np.percentile(image, percentile_value, axis=(0, 1)))
    
    return whitebalanced


def denoise(img):
    #ycbcr_img=color.rgb2ydbdr(img)
    #ycbcr_img[:, :, 0]=filters.median(ycbcr_img[:, :, 0])
    #ycbcr_img=denoise_tv_chambolle(ycbcr_img,weight=0.012, channel_axis=-1)
    #denoised_img=color.ydbdr2rgb(ycbcr_img)
    denoised_img=denoise_tv_chambolle(img,weight=0.012, channel_axis=-1)
    return denoised_img

def sk_enhance(input,thrshld=65):
    #print(input.shape)
    is_light = np.mean(input) < thrshld
    #print('np.mean(input)',np.mean(input))
    if is_light:
        #print('Low light',is_light)
        hsv_image = color.rgb2ydbdr(input)
        hsv_image[:, :, 0] = equalize_adapthist(hsv_image[:, :, 0])#,nbins=140
        hsv_adjusted = color.ydbdr2rgb(hsv_image)
        denoised_img=denoise(hsv_adjusted)
        return denoised_img
    
    else:
        return input
    
def sk_enhanceV1(input,thrshld=65):
    #print(input.shape)
    is_light = np.mean(input) < thrshld
    #print('np.mean(input)',np.mean(input))
    if is_light:
        #print('Low light',is_light)
        hsv_image = color.rgb2hsv(input)
        hsv_image[:, :, 2] = equalize_adapthist(hsv_image[:, :, 2])#,nbins=140
        hsv_adjusted = color.hsv2rgb(hsv_image)
        denoised_img=denoise(hsv_adjusted)
        return denoised_img
    
    else:
        return input
    
def sk_enhanceV3(input,thrshld=65):
    print('input shape',input.shape)
    is_light = np.mean(input) < thrshld
    #print('np.mean(input)',np.mean(input))
    if is_light:
        #print('Low light',is_light)
        
        hsv_image = color.rgb2hsv(input)
        #hsv_image[:, :, 0]=filters.median(hsv_image[:, :, 0])
        #hsv_image[:, :, 1]=filters.median(hsv_image[:, :, 1])
        hsv_image[:, :, 2] = equalize_adapthist(hsv_image[:, :, 2],clip_limit=0.0073)#,nbins=140
        
        hsv_adjusted = color.hsv2rgb(hsv_image)
        denoised_img=denoise(hsv_adjusted)
        denoised_img=percentile_whitebalance(denoised_img,97.5)
        return denoised_img
    
    else:
        return input
    
def sk_preEnhance(input,thrshld=65):
    is_light = np.mean(input) < thrshld
    if is_light:
        
        hsv_image = color.rgb2hsv(input)
        #hsv_image[:, :, 2] = equalize_adapthist(hsv_image[:, :, 2],clip_limit=0.0073)#,nbins=140
        hsv_image[:, :, 2] = equalize_adapthist(hsv_image[:, :, 2])#,nbins=140
        hsv_adjusted = color.hsv2rgb(hsv_image)
        denoised_img=denoise(hsv_adjusted)
        denoised_img=percentile_whitebalance(denoised_img,97.5)
        return ski.util.img_as_ubyte(denoised_img)
    
    else:
        return None